{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import argparse\n",
    "import copy\n",
    "import json\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "import sys\n",
    "import re\n",
    "import numpy as np\n",
    "import openai\n",
    "from typing import List, Dict, Any\n",
    "import ast\n",
    "from scan_utils import *\n",
    "from tqdm import tqdm\n",
    "\n",
    "#GPT_MODEL = \"gpt-4-turbo-preview\"  # [gpt-4-turbo-preview]\n",
    "GPT_MODEL = \"gpt-3.5-turbo\"  # [gpt-3.5-turbo-preview]\n",
    "\n",
    "setOpenAi(keyid =4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_instruction(instruction):\n",
    "    parts = instruction.split('OUT:')\n",
    "    question = parts[0].replace('IN:', '').strip()\n",
    "    answer = parts[1].strip()\n",
    "    return question, answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = []\n",
    "solutions = []\n",
    "count = 0\n",
    "N =100\n",
    "with open(\"/Users/natehu/Desktop/Tsinghua Research/TaDe/v1/scan_tasks.txt\", 'r', encoding= 'utf-8') as file:\n",
    "    for line in file:            \n",
    "            question, actions = split_instruction(line.strip())            \n",
    "            tasks.append(question)\n",
    "            actions = [action.replace(\"I_\", \"\") for action in actions.split()]\n",
    "            solutions.append(actions)\n",
    "            if count == N:\n",
    "                break\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(question):\n",
    "    prompt_cot = \"\"\" \n",
    "    Instruction:\n",
    "    Your task is to process a series of command instructions and generate the appropriate sequence of commands.\n",
    "    \n",
    "    Here are some examples:\n",
    "    Q: \"turn left\"\n",
    "    A: \"turn left\" outputs \"TURN LEFT\".\n",
    "\n",
    "    Q: \"turn right\"\n",
    "    A: \"turn right\" outputs \"TURN RIGHT\".\n",
    "\n",
    "    Q: \"jump left\"\n",
    "    A: The output of “jump left” concatenates: the output of “turn left”, the output of “jump”. “turn left” outputs “TURN LEFT”. “jump” outputs “JUMP”. So concatenating the output of “turn left” and the output of “jump” leads to “TURN LEFT” + “JUMP”. So the output of “jump left” is “TURN LEFT” + “JUMP”.\n",
    "\n",
    "    Q: \"run right\"\n",
    "    A: The output of \"run right\" concatenates: the output of \"turn right\", the output of \"run\". \"turn right\" outputs \"TURN RIGHT\". \"run\" outputs \"RUN\". So concatenating the output of \"turn right\" and the output of \"run\" leads to \"TURN RIGHT\" + \"RUN\". So the output of \"run right\" is \"TURN RIGHT\" + \"RUN\".\n",
    "\n",
    "    Q: \"look twice\"\n",
    "    A: The output of \"look twice\" concatenates: the output of \"look\", the output of \"look\". \"look\" outputs \"LOOK\". So repeating the output of \"look\" two times leads to \"LOOK\" * 2. So the output of \"look twice\" is \"LOOK\" * 2.\n",
    "\n",
    "    Q: \"run and look twice\"\n",
    "    A: The output of \"run and look twice\" concate+nates: the output of \"run\", the output of \"look twice\". \"run\" outputs \"RUN\". \"look twice\" outputs \"LOOK\" * 2. So concatenating the output of \"run\" and the output of \"look twice\" leads to \"RUN\" + \"LOOK\" * 2. So the output of \"run and look twice\" is \"RUN\" + \"LOOK\" * 2.\n",
    "\n",
    "    Q: \"jump right thrice\"\n",
    "    A: The output of \"jump right thrice\" concatenates: the output of \"jump right\", the output of \"jump right\", the output of \"jump right\". \"jump right\" outputs \"TURN RIGHT\" + \"JUMP\". So repeating the output of \"jump right\" three times leads to (\"TURN RIGHT\" + \"JUMP\") * 3. So the output of \"jump right thrice\" is (\"TURN RIGHT\" + \"JUMP\") * 3.\n",
    "\n",
    "    Q: \"walk after run\"\n",
    "    A: The output of \"walk after run\" concatenates: the output of \"run\", the output of \"walk\". \"run\" outputs \"RUN\". \"walk\" outputs \"WALK\". So concatenating the output of \"run\" and the output of \"walk\" leads to \"RUN\" + \"WALK\". So the output of \"walk after run\" is \"RUN\" + \"WALK\".\n",
    "\n",
    "    Q: \"turn opposite left\"\n",
    "    A: The output of \"turn opposite left\" concatenates: the output of \"turn left\", the output of \"turn left\". \"turn left\" outputs \"TURN LEFT\". So repeating the output of \"turn left\" twice leads to \"TURN LEFT\" * 2. So the output of \"turn opposite left\" is \"TURN LEFT\" * 2.\n",
    "\n",
    "    Q: \"turn around left\"\n",
    "    A: The output of \"turn around left\" concatenates: the output of \"turn left\", the output of \"turn left\", the output of \"turn left\", the output of \"turn left\". \"turn left\" outputs \"TURN LEFT\". So repeating the output of \"turn left\" four times leads to \"TURN LEFT\" * 4. So the output of \"turn around left\" is \"TURN LEFT\" * 4. Q: \"turn opposite right\" A: The output of \"turn opposite right\" concatenates: the output of \"turn right\", the output of \"turn right\". \"turn right\" outputs \"TURN RIGHT\". So repeating the output of \"turn right\" twice leads to \"TURN RIGHT\" * 2. So the output of \"turn opposite right\" is \"TURN RIGHT\" * 2.\n",
    "\n",
    "    Q: \"turn around right\"\n",
    "    A: The output of \"turn around right\" concatenates: the output of \"turn right\", the output of \"turn right\", the output of \"turn right\", the output of \"turn right\". \"turn right\" outputs \"TURN RIGHT\". So repeating the output of \"turn right\" four times leads to \"TURN RIGHT\" * 4. So the output of \"turn around right\" is \"TURN RIGHT\" * 4.\n",
    "\n",
    "    Q: \"walk opposite left\"\n",
    "    A: The output of \"walk opposite left\" concatenates: the output of \"turn opposite left\", the output of \"walk\". \"turn opposite left\" outputs \"TURN LEFT\" * 2. \"walk\" outputs \"WALK\". So concatenating the output of \"turn opposite left\" and the output of \"walk\" leads to \"TURN LEFT\" * 2 + \"WALK\". So the output of \"walk opposite left\" is \"TURN LEFT\" * 2 + \"WALK\".\n",
    "\n",
    "    Q: \"walk around left\"\n",
    "    A: The output of \"walk around left\" concatenates: the output of \"walk left\", the output of \"walk left\", the output of \"walk left\", the output of \"walk left\". \"walk left\" outputs \"TURN LEFT\" + \"WALK\". So repeating the output of \"walk around left\" four times leads to (\"TURN LEFT\" + \"WALK\") * 4. So the output of \"walk around left\" is (\"TURN LEFT\" + \"WALK\") * 4.\n",
    "\n",
    "    \n",
    "    Pay attention to the use of parentheses.\n",
    "    \"\"\"\n",
    "    \n",
    "    query_sys = {\n",
    "        \"role\":\"system\",\n",
    "        \"content\": prompt_cot\n",
    "    }\n",
    "    \n",
    "    query_user = {\n",
    "        \"role\":\"user\",\n",
    "        \"content\": f\"\"\"Let's think Step by Step, solve this question: {question}\n",
    "        The output should be a list of commands in a string. For example: \"RUN\" * 2 + \"TURN_LEFT\" + \"LOOK\" * 3.\n",
    "        \"\"\"\n",
    "    }\n",
    "    \n",
    "    query = [query_sys, query_user]\n",
    "    \n",
    "    result = askChatGPT(query, model=GPT_MODEL,temperature=1)\n",
    "    \n",
    "    print(\"******************************** sanity check ********************************\")\n",
    "    print(result)\n",
    "    \n",
    "    # 做mapping 得到一个式子\n",
    "    \n",
    "    queyr_assistant = {\n",
    "        \"role\":\"assistant\",\n",
    "        \"content\": result\n",
    "    }\n",
    "    \n",
    "    query.append(queyr_assistant)\n",
    "    \n",
    "    q_format = {\n",
    "        \"role\":\"user\",\n",
    "        \"content\": \"\"\"\n",
    "        What is the output? \n",
    "        Make sure only give the answer. Don't include any other information.\n",
    "        If the answer is \"RUN\" * 2 + \"TURN_LEFT\" + \"LOOK\" * 3, you should only write 'RUN' * 2 + 'TURN_LEFT' + 'LOOK' * 3.\n",
    "        \"\"\"\n",
    "    }\n",
    "    \n",
    "    query.append(q_format)\n",
    "    \n",
    "    ## 再问一次确认答案\n",
    "    \n",
    "    finalanswer = askChatGPT(query, model=GPT_MODEL,temperature=1)\n",
    "    \n",
    "    print(\"******************************** final answer ********************************\")\n",
    "    print(finalanswer)\n",
    "    \n",
    "    answer_assistant = {\n",
    "        \"role\":\"assistant\",\n",
    "        \"content\": finalanswer\n",
    "    }\n",
    "    \n",
    "    query.append(answer_assistant)\n",
    "    \n",
    "    ## 最后把这个公式给转化成一系列的命令\n",
    "    \n",
    "    rewrite_system = {\"role\": \"system\", \"content\": f\"\"\"\n",
    "    Now I have a pseudo action sequence expression with parentheses and multiplication. I need you to help me convert this into a sequence of actions without an operator sign.\n",
    "    6 examples are as follows:    \n",
    "        \n",
    "    Q: \"JUMP\" * 3\n",
    "    Rewrite: \"JUMP\" * 3\n",
    "    A: 1 JUMP 2 JUMP 3 JUMP\n",
    "\n",
    "    Q: \"RUN\" * 4 * 2\n",
    "    Rewrite: \"RUN\" * 8\n",
    "    A: 1 RUN 2 RUN 3 RUN 4 RUN 5 RUN 6 RUN 7 RUN 8 RUN\n",
    "\n",
    "    Q: \"TURN RIGHT\" + \"WALK\"\n",
    "    Rewrite: \"TURN RIGHT\" + \"WALK\"\n",
    "    A: TURN RIGHT WALK\n",
    "\n",
    "    Q: (\"TURN LEFT\" + \"LOOK\") * 2 + \"TURN LEFT\" + \"LOOK\"\n",
    "    Rewrite: (\"TURN LEFT\" + \"LOOK\") * 2 + \"TURN LEFT\" + \"LOOK\"\n",
    "    A: 1 (TURN LEFT LOOK) 2 (TURN LEFT LOOK) TURN LEFT LOOK\n",
    "\n",
    "    Q: (\"TURN RIGHT\" * 2 + \"JUMP\") * 4\n",
    "    Rewrite: (\"TURN RIGHT\" * 2 + \"JUMP\") * 4\n",
    "    A: 1 (1 TURN RIGHT 2 TURN RIGHT JUMP) 2 (1 TURN RIGHT 2 TURN RIGHT JUMP) 3 (1 TURN RIGHT 2 TURN RIGHT JUMP) 4 (1 TURN RIGHT 2 TURN RIGHT JUMP)\n",
    "\n",
    "    Q: \"TURN LEFT\" * 2 + (\"TURN RIGHT\" + \"WALK\") * 4 * 2\n",
    "    Rewrite: \"TURN LEFT\" * 2 + (\"TURN RIGHT\" + \"WALK\") * 8\n",
    "    A: 1 TURN LEFT 2 TURN LEFT 1 (TURN RIGHT WALK) 2 (TURN RIGHT WALK) 3 (TURN RIGHT WALK) 4 (TURN RIGHT WALK) 5 (TURN RIGHT WALK) 6 (TURN RIGHT WALK) 7 (TURN RIGHT WALK) 8 (TURN RIGHT WALK)\n",
    "    \"\"\"}\n",
    "    \n",
    "    Q_change = {\"role\": \"user\", \"content\": f\"\"\"The pseudo action sequence to be converted is as follows: {finalanswer} Please change it to the action sequences.\n",
    "Please JUST answer the result.\"\"\"}\n",
    "    \n",
    "    Q_now = [rewrite_system, Q_change]\n",
    "    \n",
    "    actions = askChatGPT(Q_now, model=GPT_MODEL, temperature = 1)\n",
    "\n",
    "    # 删除数字\n",
    "    text_no_numbers_and_brackets = re.sub(r'[\\d\\(\\)]', '', actions)\n",
    "    # 替换 TURN 后面的空格为下划线\n",
    "    text_formatted = re.sub(r'TURN\\s+', 'TURN_', text_no_numbers_and_brackets)\n",
    "    # 去掉多余的空格\n",
    "    text_formatted = re.sub(r'\\s+', ' ', text_formatted).strip()\n",
    "    \n",
    "    return text_formatted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "walk around left after run around left thrice\n",
      "['TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'RUN', 'TURN_LEFT', 'WALK', 'TURN_LEFT', 'WALK', 'TURN_LEFT', 'WALK', 'TURN_LEFT', 'WALK']\n"
     ]
    }
   ],
   "source": [
    "i = 78\n",
    "\n",
    "question = tasks[i]\n",
    "golden_answer = solutions[i]\n",
    "\n",
    "print(question)\n",
    "print(golden_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "******************************** sanity check ********************************\n",
      "First, let's break down the given command into smaller parts:\n",
      "\n",
      "1. \"walk around left\"\n",
      "    - This command translates to (\"TURN LEFT\" + \"WALK\") * 4 as per the provided instructions.\n",
      "    - So, the output of \"walk around left\" is (\"TURN LEFT\" + \"WALK\") * 4.\n",
      "\n",
      "2. \"run around left thrice\"\n",
      "    - This command consists of two parts:\n",
      "        a. \"run around left\"\n",
      "            - This translates to (\"TURN LEFT\" + \"RUN\") * 4.\n",
      "        b. \"thrice\"\n",
      "            - This indicates repeating the previous command (\"TURN LEFT\" + \"RUN\") * 4 three times.\n",
      "\n",
      "Now, we combine the above two parts:\n",
      "\n",
      "    Output: ((\"TURN LEFT\" + \"WALK\") * 4) + ((\"TURN LEFT\" + \"RUN\") * 4) * 3\n",
      "\n",
      "Therefore, the final answer is:\n",
      "\n",
      "    ((\"TURN LEFT\" + \"WALK\") * 4) + ((\"TURN LEFT\" + \"RUN\") * 4) * 3\n",
      "******************************** final answer ********************************\n",
      "(('TURN LEFT' + 'WALK') * 4) + (('TURN LEFT' + 'RUN') * 4) * 3\n"
     ]
    }
   ],
   "source": [
    "solution = run(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT WALK TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT RUN TURN_LEFT RUN'"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "actionList = solution.split()\n",
    "actionList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'RUN',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK',\n",
       " 'TURN_LEFT',\n",
       " 'WALK']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "golden_answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_price(number_of_token, type, model):\n",
    "    if model == \"gpt-3.5-turbo\":\n",
    "        if type == \"input\":\n",
    "            per_token = 0.5 / 1000000\n",
    "            return number_of_token * per_token\n",
    "        elif type == \"output\":\n",
    "            per_token = 1.5 / 1000000\n",
    "            return number_of_token * per_token\n",
    "    elif model == \"gpt-4-turbo-preview\":\n",
    "        if type == \"input\":\n",
    "            per_token = 10 / 1000000\n",
    "            return number_of_token * per_token\n",
    "        elif type == \"output\":\n",
    "            per_token = 30 / 1000000\n",
    "            return number_of_token * per_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0821634999999998"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_price(110494, \"output\", \"gpt-3.5-turbo\") + calculate_price(1832845, \"input\", \"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1943339"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "110494 + 1832845"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "67154"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "52543 + 14611"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "413663"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "96248 + 317415"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5586800000000001"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_price(11588, \"output\", \"gpt-4-turbo-preview\") + calculate_price(21104, \"input\", \"gpt-4-turbo-preview\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32692"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "11588 + 21104"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100170"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "31062 + 69108"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "165785"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "115643+50142\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9793693375224695"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "22030925/22495012"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4715138.12"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "0.979 * 4816280"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = {\n",
    "    \"role\":\"user\",\n",
    "    \"content\": \"Let 1≤r≤n and consider all subsets of r elements of the set {1, 2, ..., n}. What is the arithmetic mean of the smallest element of these subsets, as an expression of n and r?\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GPTlogprob(query, model):\n",
    "    response = openai.ChatCompletion.create(\n",
    "      model=model,\n",
    "      messages=query,\n",
    "      logprobs=True\n",
    "    )\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.ChatCompletion.create(\n",
    "    model=GPT_MODEL,\n",
    "    messages=[test],\n",
    "    logprobs = True\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Log Probability: 76.15\n"
     ]
    }
   ],
   "source": [
    "def calculate_average_logprob(response):\n",
    "    \"\"\"\n",
    "    Calculate the average log probability of tokens from an OpenAI API response.\n",
    "\n",
    "    Parameters:\n",
    "    response (dict): The response dictionary from OpenAI API containing 'choices' with 'logprobs'.\n",
    "\n",
    "    Returns:\n",
    "    float: The average log probability of all tokens in the response.\n",
    "    \"\"\"\n",
    "    # Extract the response's first choice (assuming the relevant data is there)\n",
    "    choice = response['choices'][0]\n",
    "\n",
    "    # Check if 'logprobs' and 'content' are available in the response\n",
    "    if 'logprobs' in choice and 'content' in choice['logprobs']:\n",
    "        # Get the list of log probabilities\n",
    "        logprobs_list = choice['logprobs']['content']\n",
    "\n",
    "        # Extract log probabilities\n",
    "        logprobs = [item['logprob'] for item in logprobs_list]\n",
    "\n",
    "        # Calculate the average log probability\n",
    "        if logprobs:\n",
    "            average_logprob = sum(logprobs) / len(logprobs)\n",
    "            return np.round(np.exp(average_logprob)*100,2)\n",
    "        else:\n",
    "            return 0  # Return 0 if no log probabilities are found\n",
    "    else:\n",
    "        raise ValueError(\"Log probabilities are missing from the response\")\n",
    "\n",
    "\n",
    "print(\"Average Log Probability:\", calculate_average_logprob(response))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "88.0"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.round(np.exp(calculate_average_logprob(response))*100,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"To find the arithmetic mean of the smallest element of all subsets of r elements of the set {1, 2, ..., n}, we need to consider how many subsets contain a particular element as the smallest element.\\n\\nLet's consider the element 1. In order for 1 to be the smallest element of a subset of size r, the remaining r-1 elements must be chosen from the set {2, 3, ..., n}. There are (n-1 choose r-1) such subsets.\\n\\nSimilarly, if we consider the element 2, in order for 2 to be the smallest element of a subset of size r, the remaining r-1 elements must be chosen from the set {3, 4, ..., n}. There are (n-2 choose r-1) such subsets.\\n\\nContinuing this pattern, if we consider the element k, in order for k to be the smallest element of a subset of size r, the remaining r-1 elements must be chosen from the set {k+1, k+2, ..., n}. There are (n-k choose r-1) such subsets.\\n\\nTherefore, the arithmetic mean of the smallest element of all subsets of r elements of the set {1, 2, ..., n} is:\\n\\n(1*(n-1 choose r-1) + 2*(n-2 choose r-1) + ... + (n-r+1)*r choose r-1) / (n choose r)\\n\\nThis expression represents the sum of the smallest elements of all subsets of r elements, divided by the total number of subsets.\""
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"To find the arithmetic mean of the smallest element of all subsets of r elements of the set {1, 2, ..., n}, we need to consider how many subsets contain a particular element as the smallest element.\\n\\nLet's consider the element 1. In order for 1 to be the smallest element of a subset of size r, the remaining r-1 elements must be chosen from the set {2, 3, ..., n}. There are (n-1 choose r-1) such subsets.\\n\\nSimilarly, if we consider the element 2, in order for 2 to be the smallest element of a subset of size r, the remaining r-1 elements must be chosen from the set {3, 4, ..., n}. There are (n-2 choose r-1) such subsets.\\n\\nContinuing this pattern, if we consider the element k, in order for k to be the smallest element of a subset of size r, the remaining r-1 elements must be chosen from the set {k+1, k+2, ..., n}. There are (n-k choose r-1) such subsets.\\n\\nTherefore, the arithmetic mean of the smallest element of all subsets of r elements of the set {1, 2, ..., n} is:\\n\\n(1*(n-1 choose r-1) + 2*(n-2 choose r-1) + ... + (n-r+1)*r choose r-1) / (n choose r)\\n\\nThis expression represents the sum of the smallest elements of all subsets of r elements, divided by the total number of subsets.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = {\n",
    "    \"role\":\"user\",\n",
    "    \"content\": \"Let 1≤r≤n and consider all subsets of r elements of the set {1, 2, ..., n}. What is the arithmetic mean of the smallest element of these subsets, as an expression of n and r?\"\n",
    "}\n",
    "query = [test]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = GPTlogprob(query, GPT_MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "76.15"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_average_logprob(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"To find the arithmetic mean of the smallest element of all subsets of r elements of the set {1, 2, ..., n}, we need to consider how many subsets will contain a particular number as the smallest element.\\n\\nLet's consider the number k (1 ≤ k ≤ n). The number k will be the smallest element of subsets that contain the number k along with (r-1) other elements chosen from the set {1, 2, ..., k-1}. The number of ways to choose (r-1) elements from {1, 2, ..., k-1} is denoted by C(k-1, r-1).\\n\\nTherefore, the number k will be the smallest element in C(k-1, r-1) subsets. Since k can range from 1 to n, we need to sum up the number of subsets for each k and divide by the total number of subsets, which is C(n, r).\\n\\nThe arithmetic mean of the smallest element can be calculated as follows:\\n\\nArithmetic Mean = Σ[k * C(k-1, r-1)] / C(n, r)\\n                  = Σ[(k * C(k-1, r-1))] / C(n, r) for k = 1 to n\\n\\nTherefore, the arithmetic mean of the smallest element of subsets of r elements of the set {1, 2, ..., n} is given by the expression:\\n\\nΣ[(k * C(k-1, r-1))] / C(n, r) where k ranges from 1 to n.\""
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user',\n",
       "  'content': 'Let 1≤r≤n and consider all subsets of r elements of the set {1, 2, ..., n}. What is the arithmetic mean of the smallest element of these subsets, as an expression of n and r?'}]"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "query.append({'role': 'assistant', 'content': response.choices[0].message.content})\n",
    "query.append({'role': 'user', 'content': \"therefore, what is the final answer\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_final = GPTlogprob(query, GPT_MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "76.56"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_average_logprob(response_final)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
